$ANSI ||= false

class String

  class << self

    attr_accessor :ansi_codes

    def ansi_codes
      @ansi_codes ||= {}
    end

    def define_ansi(definitions = {})
      self.ansi_codes = ansi_codes.merge(definitions)
      define_shortcuts(definitions)
    end

    def define_shortcuts(definitions = {})
      definitions.keys.each do |meth|
        unless instance_methods.include?(meth.to_s)
          define_method(meth) { with_ansi(meth) }
        end
      end
    end

  end

  define_ansi :normal         => "0m",      :bold           => "1m"
  define_ansi :underline      => "4m",      :blink          => "5m"
  define_ansi :reverse_video  => "7m",      :invisible      => "8m"
  define_ansi :black          => "30m",     :red            => "31m"
  define_ansi :green          => "32m",     :yellow         => "33m"
  define_ansi :blue           => "34m",     :magenta        => "35m"
  define_ansi :cyan           => "36m",     :white          => "37m"
  define_ansi :black_bg       => "40m",     :red_bg         => "41m"
  define_ansi :green_bg       => "42m",     :yellow_bg      => "43m"
  define_ansi :blue_bg        => "44m",     :magenta_bg     => "45m"
  define_ansi :cyan_bg        => "46m",     :white_bg       => "47m"
  define_ansi :clear_line     => "K",       :clear_screen   => "2J"
  define_ansi :go_home        => "0H",      :go_to_end      => "80L"

  # Wrap a string with an arbitrary ansi code and the ansi normal code
  def with_ansi(*codes)
    use_ansi? ? "#{sym_to_ansi(*codes)}#{self}#{sym_to_ansi(:normal)}" : self
  end

  # Just a little metaprogramming shortcut
  def ansi_codes; self.class.ansi_codes; end

  private
  def sym_to_ansi(*symbols)
    symbols.inject("") do |string, symbol|
      if code = ansi_codes[symbol]
        string << (code.is_a?(Array) ? sym_to_ansi(*code) : "\e[#{code}")
      end
      string
    end
  end
  
  # determine whether we have ansi support or ANSI enabled
  def use_ansi?
    $ANSI && RUBY_PLATFORM !~ /win32/i
  end

end
